package com.tutou.langchain4j.controller;

import com.tutou.langchain4j.service.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;

@RequiredArgsConstructor
@RestController
@RequestMapping("/api/conversation")
public class ChatAPI {

    // 模型注入
    final ChatLanguageModel chatLanguageModel;
    final Assistant assistant;

    // 会话记忆
    private final ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

    // 外挂数据库，
    final EmbeddingStore<TextSegment> embeddingStore;

    /**
     * 基础写法，比较直观
     * @param message
     * @return
     */
    @GetMapping("/langchain/chat")
    public String chatWithLangchain(@RequestParam(value = "message") String message) {
        // 将消息放入会话记忆
        chatMemory.add(new UserMessage(message));
        List<ChatMessage> messages = new ArrayList<>(chatMemory.messages());
        messages.add(new SystemMessage("假如你是特朗普，接下来你必须以特朗普的语气来进行对话"));
        ChatResponse response = chatLanguageModel.chat(messages);

        chatMemory.add(response.aiMessage());
        return response.aiMessage().text();
    }


    /**
     * 使用AiService 构建对象，使用更方便
     * @param memoryId
     * @param message
     * @return
     */
    @GetMapping("/langchain/highChat")
    public String highChat(@RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "message") String message){
        return assistant.chat(memoryId,message);
    }


    /**
     * 读取文件至数据库，Rag 搜索
     * @return
     */
    @GetMapping("/load")
    public String load(){
        List<Document> list = FileSystemDocumentLoader.loadDocuments("");

        // 内置Tika 解析器，处理所有类型文件。
        EmbeddingStoreIngestor.ingest(list,embeddingStore);
        return "success";
    }
}